{-# LANGUAGE RankNTypes,PatternGuards,NoMonomorphismRestriction,FlexibleInstances,UnboxedTuples #-}
{-# LANGUAGE ViewPatterns,DeriveGeneric,ScopedTypeVariables,BangPatterns,FlexibleContexts #-}


module Utils where

import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import Data.HashSet ( HashSet )
import qualified Data.HashSet as HS
import Data.IntSet ( IntSet )
import qualified Data.IntSet as IS
import Data.Maybe  
import Data.List 
import qualified Data.Text as T (unpack)

import LLVM.Analysis


type ValueIds = IntSet

refs :: IsValue a => a -> [String]
refs = mapMaybe toVarName . HS.toList . refVals

refVals :: IsValue a => a -> HashSet Value
refVals = HS.filter (not . isConstant). valueRefs

valueRefs :: IsValue a => a -> HashSet Value
valueRefs = go HS.empty . HS.singleton . toValue        
  where    
    go ::  HashSet Value -> HashSet Value -> HashSet Value
    go visited q
      | HS.null vals = visited
      | otherwise =
        let visited' = visited `HS.union` vals
            q' = foldl' addValuesFrom HS.empty (HS.toList vals)
        in go visited' q'
      where
        vals = HS.difference q visited
    ---
    addValuesFrom :: HashSet Value -> Value -> HashSet Value
    addValuesFrom q v =
      case valueContent' v of
        InstructionC i@StoreInst {storeValue = sv} -> HS.insert sv q
        InstructionC i@PhiNode {} -> foldl' (flip HS.insert) q vs
          where vs = instructionOperands i ++ (maybeToList $ phi2Var v)
        InstructionC i -> foldl' (flip HS.insert) q (instructionOperands i)
        FunctionC f -> foldl' (flip HS.insert) q (funcExitValue f)
        _ -> HS.insert (stripBitcasts v) q             
{-# INLINE valueRefs #-} 

-- | the possible variable referred by a PHI instruction
phi2Var :: IsValue a => a -> Maybe Value
phi2Var pv = case valueContent' (toValue pv) of 
   InstructionC PhiNode {phiIncomingValues = pvs} ->
     case partition (isConstant . fst) pvs of 
       ([],[]) -> Nothing
       (vs,[]) -> listToMaybe2 $ mapMaybe getVar vs
       (_, vs) -> listToMaybe2 $ mapMaybe getVar vs 
   _ -> Nothing
   where
      getVar (val,lab) = let ls = lkpVars val lab in
        if not (null ls) then listToMaybe2 ls
        else case valueContent' val of
          InstructionC LoadInst {loadAddress = la} -> Just la
          _ -> Nothing 
      lkpVars v = mapMaybe (matchVar v). basicBlockInstructions. toBB 
      listToMaybe2 ls = case nub ls of
          [v] -> Just v   
          _   -> Nothing  
      matchVar v i = case i of 
        StoreInst {storeValue = sv, storeAddress = sa} ->
           if sv == v then Just sa else Nothing
        _ -> Nothing
      toBB v = case valueContent v of
        BasicBlockC bb -> bb
        _ -> error ("ParaSlicer.phi2Var: Expected basic block: " ++ show v)
{-# INLINE phi2Var #-} 
        
toValName, toVarName' :: IsValue a => a -> String
toValName v = maybe (valStr v) show (valueName v)
  where valStr v = "_" ++ show (valueUniqueId v)

toVarName' val = fromMaybe (toValName val) . toVarName $ toValue val
 
toVarName :: Value -> Maybe String
toVarName val =
  case valueContent' val of
    GlobalVariableC gv -> Just $ show (globalVariableName gv)
    InstructionC i@AllocaInst {} -> Just $ instName i    
    ArgumentC av -> Just $ argName av
--    InstructionC i@PhiNode {} -> liftM show (valueName i)    -- for phi vars
    _ -> Nothing   
  where  
    instName i = maybe (instIdStr i) show (instructionName i) ++ 
                 maybe "@**" (show . functionName) (instructionFunction i)
    instIdStr i = "_" ++ show (instructionUniqueId i)
    argName av = show (argumentName av) ++ (show . functionName $ argumentFunction av)


-----
---
findVals :: IntMap Value -> IntSet -> [Value]
findVals m ks = IM.elems (IM.filterWithKey keyF m)
  where keyF k _ = IS.member k ks  


genValueMap :: Module -> IntMap Value
genValueMap m = IM.fromList (zip allIds allVals)
  where   
    allVals = allValues m
    allIds = map valueUniqueId allVals


--------
valueLine :: IsValue a => a -> [Int] 
valueLine = mapMaybe metadataToLine . valueMetadata

metadataToLine :: Metadata -> Maybe Int
metadataToLine md@MetaSourceLocation { } = Just . fromIntegral . metaSourceRow $ md
metadataToLine md@MetaDWSubprogram {} = Just . fromIntegral . metaSubprogramLine $ md
metadataToLine md@MetaDWLexicalBlock { } = Just . fromIntegral . metaLexicalBlockRow $ md
metadataToLine md@MetaDWVariable {} = Just . fromIntegral . metaGlobalVarLine $ md
metadataToLine md@MetaDWBaseType {} = Just . fromIntegral . metaBaseTypeLine $ md
metadataToLine md@MetaDWDerivedType {} = Just . fromIntegral . metaDerivedTypeLine $ md
metadataToLine md@MetaDWCompositeType {} = Just . fromIntegral . metaCompositeTypeLine $ md
metadataToLine md@MetaDWLocal {} = Just . fromIntegral . metaLocalLine $ md
metadataToLine md@MetaDWNamespace {} = Just . fromIntegral . metaNamespaceLine $ md
metadataToLine md@MetaDWTemplateTypeParameter {} = 
       Just . fromIntegral . metaTemplateTypeParameterLine $ md
metadataToLine md@MetaDWTemplateValueParameter {} = 
       Just . fromIntegral . metaTemplateValueParameterLine $ md
metadataToLine _ = Nothing

--------
valueFunction :: IsValue a => a -> Maybe Function
valueFunction v = case valueContent v of
    FunctionC f -> Just f
    ArgumentC a -> Just (argumentFunction a)
    BasicBlockC b -> Just (basicBlockFunction b)
    InstructionC i -> instructionFunction i
    _ -> Nothing
   
-------------------------------- 
-----------------------------
-- stealing from LLVM.Analysis
allValues :: Module -> [Value]
allValues m = allVals
  where
    fs = moduleDefinedFunctions m
    allArgs = concatMap functionParameters fs
    allBlocks = concatMap functionBody fs
    allInsts = concatMap basicBlockInstructions allBlocks
    allVals = concat [ map toValue fs
                     , map toValue (moduleGlobalVariables m)
                     , map toValue (moduleExternalValues m)
                     , map toValue (moduleExternalFunctions m)
                     , map toValue (moduleAliases m)
                     , map toValue allBlocks
                     , map toValue allInsts
                     , map toValue allArgs
                     ]

funcExitValue :: Function -> [Value]
funcExitValue f = 
   case functionExitInstruction f of
      Just i@RetInst {retInstValue = Just rv} -> [toValue i, rv]
      _ -> []

funcAllocInsts :: Function -> [Instruction]
funcAllocInsts f = allocInsts
  where
--    allVars = zip allocVars (repeat $ identifierAsString (functionName f))
    allocInsts = filter isAllocaInst $ functionInstructions f
--    allocVars =  map identifierAsString $ mapMaybe instructionName allocInsts  
 
isAllocaInst :: Instruction -> Bool
isAllocaInst AllocaInst {} = True
isAllocaInst _ = False

----
isAggregate :: (IsValue v) => v -> Bool              
isAggregate v = isGetElem v || isAggType v

isGetElem, isAggType :: (IsValue v) => v -> Bool
isGetElem v =  case valueContent' v of 
     InstructionC GetElementPtrInst {} -> True
     _ -> False
    
isAggType v =  case valueType v of
     TypeArray _ _ -> True
     TypeStruct _ _ _ -> True
     _ -> False    
  
isMemcpy :: Value -> Bool
isMemcpy v =
  case valueContent' v of
    ExternalFunctionC ExternalFunction { externalFunctionName = fname } ->
      show fname == "@llvm.memcpy.p0i8.p0i8.i32" || show fname == "@llvm.memcpy.p0i8.p0i8.i64"
    _ -> False   

isConstant :: Value -> Bool
isConstant v = case valueContent' v of
  ConstantC _ -> True
  _ -> False
  
------
valueSrcLn :: IsValue a => a -> Maybe (FilePath, [Int])  
valueSrcLn v = do
  func <- valueFunction v
  let mds = filter isSubprogramMetadata $ functionMetadata func
      line = valueLine v
  case mds of
    [md] -> do
      ctxt <- metaSubprogramContext md
      let f = metaFileSourceFile ctxt
          d = metaFileSourceDir ctxt
          absSrcFile = T.unpack f  -- (T.unpack d) </> (T.unpack f)
      return (absSrcFile, line)  
    _ -> Nothing
  where
    isSubprogramMetadata :: Metadata -> Bool
    isSubprogramMetadata MetaDWSubprogram {} = True
    isSubprogramMetadata _ = False